from visual import *
import ode
from decimal import *

#getcontext().prec = 12

class QuadMotor(object):
    def __init__(self, motor_frame, body, space, motor_pos, motor_rad, motor_dir, color):
        self.body = body
        self.motor_frame = motor_frame
        self.vis_motor = cylinder(frame = motor_frame, pos = motor_pos, axis = (0, 1, 0), radius = motor_rad )
        self.vis_motor.color = color
        shaft_pos = self.tuple_to_vec(motor_pos) + vector(0, 1, 0)
        
        self.shaft = cylinder(frame = motor_frame, pos = shaft_pos , axis = (0, 1, 0), radius = 0.2)
        self.motor_mass = ode.Mass()
        self.motor_mass.setCylinder(20, 2, motor_rad, 1)
        mass_pos = motor_pos
        mass_pos = (motor_pos[0],0,motor_pos[2])
        
        self.motor_mass.translate(mass_pos)
        
        self.motor_geom = ode.GeomCylinder(radius = motor_rad, length = 1)
        self.motor_geom.setPosition(motor_pos)
        self.motor_geom.setRotation([1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0])
        self.motor_transform = ode.GeomTransform(space)
        self.motor_transform.setGeom(self.motor_geom)
        self.motor_transform.setBody(body)
        
        self.shaft_geom = ode.GeomCylinder(radius = 0.2, length = 1)
        self.shaft_geom.setPosition(shaft_pos)
        self.shaft_geom.setRotation([1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0])
        self.shaft_transform = ode.GeomTransform(space)
        self.shaft_transform.setGeom(self.shaft_geom)
        self.shaft_transform.setBody(body)
    
        self.propeller = None
        self.motor_joint = None
        
        self.motor_dir = motor_dir
        
        self.motor_speed = 0.0
        self.motor_rel_speed = 0.0
        self.motor_const = Decimal(5)
        
    def tuple_to_vec(self, tup ):
        return vector(tup[0],tup[1],tup[2])
    
    def get_propeller(self):
        return self.propeller
    
    def get_mass(self):
        return self.motor_mass
    
    def get_pos(self):
        a = self.vis_motor.pos
        return a
    
    def attach_propeller(self, world, propeller):
        self.propeller = propeller

        self.motor_joint = ode.HingeJoint(world)
        
        self.motor_joint.setAnchor(self.get_pos())
        self.motor_joint.attach(propeller.get_body(), self.body)
        self.motor_joint.setAxis((0,1,0))
        
        self.set_speed(0)
        
        self.motor_joint.setParam(ode.ParamVel, (self.motor_speed * 30) * self.motor_dir)
        self.motor_joint.setParam(ode.ParamFMax, 22)
       
    def get_propeller(self):
        return self.propeller
    
    def set_speed(self, motor_speed):
        speed_limit = 10
        self.motor_speed = max(min(motor_speed, speed_limit), 0)
        self.motor_joint.setParam(ode.ParamVel, (self.motor_speed * 40) * Decimal(str(self.motor_dir)))
        
    def set_dir(self, dir):
        self.motor_dir = dir
        
    def get_dir(self):
        return self.motor_dir
    
    def get_speed(self):
        return self.motor_speed
    
    def get_rel_speed(self):
        return self.motor_rel_speed + 0.0
    
    def set_rel_speed(self, rel_speed):
        self.motor_rel_speed = rel_speed + 0.0
        self.set_speed(self.motor_speed)
        
    def update(self):
        
        self.body.addTorque(self.motor_frame.up * (pow(self.motor_speed, 2)* 50 * self.motor_dir))
        force = Decimal(str(pow(self.motor_speed, 2)))
        thrust = force * self.motor_speed
        self.body.addForceAtPos(self.motor_frame.up * thrust ,self.get_pos())
        if self.propeller is not None:
            self.propeller.update()
        
